#! /usr/bin/env python
# -*- coding: utf8 -*-

# Alexandre Coster <acoster@rhrk.uni-kl.de>
# AG DBIS - Technische Universität Kaiserslautern

"""
This file contains the three classes that manage transformations: L{root}, L{transform} and L{attribute}.
"""

import sys
import copy
import math
import random

class root:
    """
    This class implements an callable object which abstracts the root transformation.
    """

    def __init__(self, node, probDistr):
        """
        Constructor for the C{root} class.

        @parameter node: C{<transformation>} node.
        @type node: Element
        @parameter probDistr: Dictionary containing L{probability distributions<mxml.prob.probDistr>}. The keys are the distributions
        names.
        @type probDistr: dict of L{probDistr<mxml.prob.probDistr>}
        """

        self.__copyRate     = float(node.find("copyRate").text.strip())  / 100.0
        self.__untouched    = float(node.find("untouched").text.strip()) / 100.0

        if node.find("selectProbDistr") != None:
            self.__selectProbDistr = probDistr[node.find("selectProbDistr").text]
            self.__errorRate = None
        else:
            self.__selectProbDistr = None
            self.__errorRate = float(node.find("errorRate").text.strip()) / 100.

        if node.find("dupRate") != None :
            self.__dupRate   = float(node.find("dupRate").text.strip()) / 100.0
            self.__dupDistro = probDistr[node.find("dupDistr").text]
        else:
             self.__dupRate = None

        if node.find("errorByWords") != None:
            self.__byWords = True
            self.__byChars = False
        elif node.find("errorByChars") != None:
            self.__byWords = False
            self.__byChars = True
        else:
            self.__byWords = self.__byChars = False

    def __call__(self, nodes=[], transforms = {}, addIDs = False, resultAmmount = 1, outFile = None):
        """
        Applies the transformation to the XML document. It will generate a quantity of new documents, as specified by the 
        C{resultAmmount} parameter.

        @parameter nodes: Nodes of the root type.
        @type nodes: list of Element
        @parameter transforms: Dictionary of L{transformations<transform>}, with tagnames as keys.
        @type transforms: dict of L{transform}
        @parameter addIDs: Whether IDs should be added to the copied and original trees.
        @type addIDs: bool
        @parameter resultAmmount: Ammount of documents that will be generated.
        @type resultAmmount: int
        @parameter outFile: File where "statistics" will be written to. None for none.
        @type outFile: file or NoneType
        @return: Copied (transformed or not) nodes. A list of list is returned, each second-level list being a transformed document.
        @rtype: list of list of Element
        """

        toReturn = []
        toCopy   = random.sample(nodes, int(self.__copyRate * len(nodes)))

        if addIDs:
            nextID = 0
            for i in toCopy:
                i.set("xxmlCopyID", "%d" % (nextID,))
                nextID += 1

        for i in range(resultAmmount):
            copies    = [copy.deepcopy(x) for x in toCopy]

            if self.__dupRate != None and int(self.__dupRate * len(copies)) > 0:
                selected = random.sample(xrange(len(copies)), int(self.__dupRate * len(copies)))
                tmpCopies = []

                for i in selected:
                    dupAmm = int(self.__dupDistro())
                    if dupAmm <= 0:
                        continue
                    [tmpCopies.append(copy.deepcopy(copies[i])) for x in xrange(dupAmm)]
                copies += tmpCopies


            toReturn.append(copies)
            victims   = random.sample(copies, int(len(copies) * (1.0 - self.__untouched)))

            for i in victims:
                allDescendants = [x for x in i.findall(".//*") if x.tag in transforms.keys() and x.text != None and len(x.text) > 0]

                if self.__selectProbDistr != None:
                    count    = len(allDescendants)
                    errors   = int(round(self.__selectProbDistr()))
                    wVictims = [(1.0/count, x) for x in allDescendants]
                elif self.__byChars:
                    count    = float(sum([len(x.text) for x in allDescendants if x.text != None]))
                    errors   = int(math.ceil(self.__errorRate * count))
                    wVictims = [(len(x.text) / count, x) for x in allDescendants if x.text != None]
                elif self.__byWords:
                    count  = float(sum([len(x.text.split(" ")) for x in allDescendants if x.text != None]))
                    errors = int(math.ceil(self.__errorRate * count))
                    wVictims = [(len(x.text.split(" ")) / count, x) for x in allDescendants if x.text != None]
                else:
                    count  = len(allDescendants)
                    errors = int(math.ceil(self.__errorRate * count))
                    wVictims = [(1.0/count, x) for x in allDescendants]

                if outFile != None:
                    outFile.write("%d %d\n" % (count, errors))

                if errors < 0:
                    errors = 0

                if len(wVictims) > 0:
                    maxWeight = max(wVictims, key=lambda x:x[0])

                for i in range(errors):
                    if len(wVictims) == 0:
                         break

                    victimNode = self.__weighted_choice(wVictims, maxWeight[0])
                    transforms[victimNode.tag](node=victimNode)
                    random.shuffle(victimNode)

        return toReturn

    def __weighted_choice(self, lista, max_val):
        while True:
            i = random.choice(lista)
            if random.random() < i[0]:
                return i[1]

class transform:
    """
    This class implements a node transformation.
    """

    def __init__(self, node, algorithms):
        """
        @parameter node: Configuration node.
        @type node: Element
        @parameter algorithms: Dictionary containing algorithms.
        @type algorithms: dict of objects that implement IAlgorithm.
        """
        self.__attributes        = {}
        self.__cdata             = None

        for i in node.findall("attribute"):
            name = i.get("name")
            self.__attributes[name] = attribute(i, algorithms)

        cdata = node.find("chars")
        if cdata != None:
            self.__attributes["mxmlCDATAattr"] = attribute(cdata, algorithms)

    def __call__(self, node):
        """
        Applies the transformation over C{node}.

        @parameter node: Node which will be transformed.
        @type node: Element
        """

        victim = random.choice(self.__attributes.keys())

        if victim == "mxmlCDATAattr":
            if node.text != None:
                result = self.__attributes[victim](node.text)
            else:
                return

            node.text = result
        else:
            if node.get(victim) != None:
                result = self.__attributes[victim](node.get(victim))
                node.set(victim, result)

class attribute:
    """
    Class that represents a transformation applied to an attribute.
    """

    def __init__(self, node, algorithms):
        """
        @parameter node: Configuration node.
        @type node: Element
        @parameter algorithms: Dictionary containing algorithms.
        @type algorithms: dict of objects that implement IAlgorithm.
        """
        self.__algorithms = []
        totalProb = 0

        for i in node.findall("useAlgo"):
            algorithm  = algorithms[i.find("name").text.strip()]
            prob       = float(i.find("useProb").text.strip())
            totalProb += prob

            self.__algorithms.append((algorithm, prob / 100.0))

        if totalProb != 100:
            print "Total prob is not 100."
            sys.exit()

    def __call__(self, string):
        """
        Transforms C{string}.

        @parameter string: String to be transformed.
        @type string: str
        """
        while True:
            for i in self.__algorithms:
                if random.random() < i[1]:
                    return i[0](string)
